'''
接入tugraph，导出schema，然后接入llm，生成cypher语句,先调试好提示词生成cypher，然后到机房跑一下
'''

from langchain_openai import ChatOpenAI
import json, textwrap,re

def llm_res(user = '乙肝（疾病）吃什么药'):
    # 1. 读文件
    with open('data/apikey.txt', 'r', encoding='utf-8') as file:
        api_key = file.read()

    with open('data/schema.json', 'r', encoding='utf-8') as f:
        schema = json.load(f)

    # 2. 转成紧凑字符串（节省 token）
    schema_str = json.dumps(schema, ensure_ascii=False, separators=(',', ':'))

    # 3. 拼 prompt

    prompt = """
    现在你是一个数仓管理员，我会给你tugraph的schema，也会给你用户问题，你需要根据用户问题生成适合tugraph的cypher语句查询。
    要求：只返回cypher语句，其他都不返回。
    schema:{schema}
    """


    res = ChatOpenAI(
    model="deepseek-ai/DeepSeek-V3",
    temperature=0,
    openai_api_key=api_key,
    openai_api_base="https://api.siliconflow.cn/v1"
    )

    from langchain_core.prompts import ChatPromptTemplate
    prompt_0 = ChatPromptTemplate.from_messages(
    [
        ("system",prompt,),
        ("user", "问题：{input}"),
    ]
    )

    chain_0 = prompt_0 | res

    res_msg =chain_0.invoke(
    {
       "input": user,
       "schema":schema
    }
    )

    print(res_msg.content)
    r = re.search(r'RETURN\s+([^;]+)', res_msg.content, flags=re.I).group(1).strip()
    from neo4j import GraphDatabase
    URI = "bolt://localhost:7687"
    AUTH = ("admin", "73@TuGraph")
    with (GraphDatabase.driver(URI, auth=AUTH) as client):
        session = client.session(database="default")
        query = res_msg.content
        # "match (n1:disease {name:'乙肝'})-[r:is_drug]-(n2) return n2.name"
        ret = session.run(query)
        result = [item.get(r) for item in ret.data()]
    print(result)

    prompt = "你现在是一个医生，我会告诉你病人的问题，你的同事已经帮你查询出病人所需要的东西，请你组织语言，回复病人。注意，如果同事给你的结果为空，请回复：对不起，这个问题我没有查到，不要输出其他东西。同事回复:{response}"

    doc = ChatOpenAI(
    model="deepseek-ai/DeepSeek-V3",
    temperature=0,
    openai_api_key=api_key,
    openai_api_base="https://api.siliconflow.cn/v1"
    )
    prompt_1 = ChatPromptTemplate.from_messages(
    [
        ("system",prompt,),
        ("user", "问题：{input}"),
    ]
    )

    chain_1 = prompt_1 | doc

    doc_msg =chain_1.invoke(
    {
       "input": user,
       "response":str(result)
    }
    )

    print(doc_msg.content)

if __name__ == '__main__':
    llm_res('头疼（症状）用什么药')